import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
from utils import to_abs_path


def flower_dataloader(batch_size=2, resize=(227, 227)):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Resize(resize),
    ])
    train = ImageFolder(root=to_abs_path('./data/flower_data/train'), transform=transform)
    test = ImageFolder(root=to_abs_path('./data/flower_data/val'), transform=transform)
    #     print(train.class_to_idx)
    #     img, label = train[0]
    #     plt.imshow(img.permute(1, 2, 0))
    #     plt.show()
    #     print(label)
    train_loader = DataLoader(train, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(test, batch_size=batch_size, shuffle=False)
    #     batch = next(iter(train_loader))
    #     images, labels = batch
    #     idx2class = {v: k for k, v in train.class_to_idx.items()}
    #
    #     fig, axs = plt.subplots(1, batch_size)
    #     for i in range(batch_size):
    #         axs[i].imshow(images[i].permute(1, 2, 0))
    #         axs[i].set_title(idx2class[labels[i].item()])
    #         axs[i].axis('off')
    #     plt.show()
    return train_loader, test_loader


def cifar10_dataloader(batch_size=2, resize=(227, 227)):
    # 训练集增强：丰富数据多样性
    train_transform = transforms.Compose([
        transforms.Resize(resize),  # 调整尺寸为 AlexNet 输入
        transforms.RandomCrop(resize, padding=32, padding_mode='reflect'),  # 随机裁剪+反射填充
        transforms.RandomHorizontalFlip(p=0.5),  # 50%概率水平翻转
        transforms.RandomRotation(15, fill=(0,)),  # 随机旋转±15度
        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),  # 颜色扰动
        transforms.RandomGrayscale(p=0.1),  # 10%概率转为灰度图
        transforms.ToTensor(),  # 转为张量
        transforms.Normalize(  # CIFAR-10 预计算的均值和标准差
            mean=[0.4914, 0.4822, 0.4465],
            std=[0.2023, 0.1994, 0.2010]
        ),
    ])

    # 测试集仅做必要变换：不增强，仅调整尺寸和归一化
    test_transform = transforms.Compose([
        transforms.Resize(resize),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.4914, 0.4822, 0.4465],
            std=[0.2023, 0.1994, 0.2010]
        ),
    ])
    train = torchvision.datasets.CIFAR10(root=to_abs_path('./data/cifar10'), train=True, download=True,
                                         transform=train_transform)
    test = torchvision.datasets.CIFAR10(root=to_abs_path('./data/cifar10'), train=False, download=True,
                                        transform=test_transform)
    train_loader = DataLoader(
        train,
        batch_size=batch_size,
        shuffle=True,
        num_workers=8,
        pin_memory=True
    )
    test_loader = DataLoader(
        test,
        batch_size=batch_size,
        shuffle=False,
        num_workers=8,
        pin_memory=True
    )
    return train_loader, test_loader


def do():
    pass


if __name__ == '__main__':
    do()
